from app.utils.cv_tools import process_image_roi
from app.utils.file_tools import clear_directory, clear_logs, extract_filename_without_extension
from app.utils.pymupdf_tools import get_pages_from_range, page_to_img, pdf_get_content_region
from app.utils.timing import timing_decorator
from tests.base_test import base_test_case

import cv2

logger = base_test_case.get_logger(__name__)
TEST_DATA_DIR = base_test_case.test_data_dir
OUTPUT_DATA_DIR = base_test_case.output_data_dir

# 全局变量存储pipeline实例
_pipeline_instances = {}

# 全局变量存储 模型缓存
_all_model_instances = {}


@timing_decorator(prefix="获取{pipeline_name}实例")
def get_pipeline_instance(pipeline_name, **init_params):
    """
    获取pipeline实例，如果不存在则创建并缓存

    Args:
        pipeline_name: pipeline名称
        **init_params: 初始化参数

    Returns:
        pipeline实例
    """
    global _pipeline_instances

    if pipeline_name not in _pipeline_instances:
        # 默认初始化参数
        default_init_params = {
            "device": "cpu",
            "initial_predictor": False
        }

        # 合并用户提供的参数
        default_init_params.update(init_params)
        from paddlex import create_pipeline
        _pipeline_instances[pipeline_name] = create_pipeline(
            pipeline=pipeline_name,
            **default_init_params
        )
        logger.info(f"实例创建完成：{pipeline_name}")

    return _pipeline_instances[pipeline_name]


@timing_decorator(prefix="通用版面分析v3")
def general_layout_analysis(input_path, pipeline_name="PP-StructureV3", prefix=None, **predict_params):
    if prefix is None:
        prefix = pipeline_name

    # 获取pipeline实例
    pipeline = get_pipeline_instance(pipeline_name)

    # 默认推理参数
    default_predict_params = {
        "input": input_path,
        "use_doc_preprocessor": False,  # 是否使用文档图像预处理模块,文档方向分类可自动识别文档的四个方向（0°、90°、180°、270°），确保文档以正确的方向进行后续处理
        "use_doc_orientation_classify": False,
        "use_doc_unwarping": False,
        "use_textline_orientation": False,  # 是否使用文本行方向分类模块
        "use_region_detection": True,
        "use_table_recognition": True,
        "use_formula_recognition": True,
        "use_chart_recognition": False,
        "use_seal_recognition": False,
        "format_block_content": True, # paddlex 3.3.0 控制是否将 block_content 中的内容格式化为Markdown格式

        # WirelessTableStructureRecognition 中的 model_name 替换为 SLANet_plus
        "use_e2e_wireless_table_rec_model": False,
        # 是否启用无线表格端到端预测模式，不启用时采用表格单元格检测模型预测结果填充至HTML表格，启用时采用端到端表格结构识别模型的单元格预测结果填充至HTML表格。二者在不同场景下效果不同，请根据实际情况选择。
        "use_wireless_table_cells_trans_to_html": False,
        # 是否启用无线表单元格直转HTML模式，不启用时正常通过表格结构识别模块预测得到表格的HTML结构，启用时直接基于无线表单元格检测模型预测的各单元格位置排布与几何关系构建HTML结构。此模式在不同场景下效果不同，请根据实际情况选择。一般来说，建议在无线表单元格检测非常准确的场景下开启此模式，以获得更加准确的端到端表格识别结果。

    }

    # 合并用户提供的参数
    default_predict_params.update(predict_params)

    # 执行预测
    output = pipeline.predict(**default_predict_params)

    for res in output:
        res.save_to_img(f"./output/{prefix}/")
        res.save_to_json(f"./output/{prefix}/")
        res.save_to_xlsx(f"./output/{prefix}/")
        res.save_to_html(f"./output/{prefix}/")
        res.save_to_markdown(f"./output/{prefix}/")


@timing_decorator(prefix="通用表格识别")
def table_recognition(input_path, pipeline_name="table_recognition", prefix=None, **predict_params):
    """
    通用表格识别函数
    Args:
        input_path: 输入文件路径或numpy数组
        pipeline_name: pipeline名称
        prefix: 输出目录前缀
        **predict_params: 预测参数
    """
    if prefix is None:
        prefix = pipeline_name

    # 获取pipeline实例
    pipeline = get_pipeline_instance(pipeline_name)
    # TODO测试说明 table_recognition 与 table_recognition_v2 在表格识别能力和效果 一般
    """表格检测
        1 探测表格坐标区域 识别正确性 平均 98%,
        2 合并使用原则:结合模型特点:可以锁定单元格 高与款与pymupdf工具配合使用,提取单元格坐标内容
        3 最好使用 table_recognition 类实例完成表格检测
    """

    # table_recognition 产线 独有推理参数   use_table_cells_ocr_results

    # table_recognition_v2 默认推理参数
    default_predict_params = {
        "input": input_path,
        "use_doc_unwarping": False,  # 是否使用文档扭曲矫正模块
        "use_layout_detection": True,  # 是否使用版面检测模块
        "use_doc_preprocessor": False,  # 是否使用文档图像预处理模块,文档方向分类可自动识别文档的四个方向（0°、90°、180°、270°），确保文档以正确的方向进行后续处理
        "use_doc_orientation_classify": False,  # 是否使用文档方向分类模块

        "use_table_orientation_classify": False,  # 是否使用表格方向分类模块
        "use_ocr_results_with_table_cells": True,
        # 是否启用单元格OCR模式，不启用时采用全局OCR结果填充至HTML表格，启用时基于单元格预测结果，对全局OCR结果基于单元格进行切分和重识别，避免出现文字跨单元格连接或缺失的情况。此模式会略微增加推理耗时，并且在不同场景下效果可能不同，请根据实际情况选择。
        "use_e2e_wired_table_rec_model": True,
        # 是否启用有线表格端到端预测模式，不启用时采用表格单元格检测模型预测结果填充至HTML表格，启用时采用端到端表格结构识别模型的单元格预测结果填充至HTML表格。二者在不同场景下效果不同，请根据实际情况选择。
        "use_wired_table_cells_trans_to_html": False,
        # 是否启用有线表单元格直转HTML模式，不启用时正常通过表格结构识别模块预测得到表格的HTML结构，启用时直接基于有线表单元格检测模型预测的各单元格位置排布与几何关系构建HTML结构。此模式在不同场景下效果不同，请根据实际情况选择。一般来说，建议在有线表单元格检测非常准确的场景下开启此模式，以获得更加准确的端到端表格识别结果。

        # WirelessTableStructureRecognition 中的 model_name 替换为 SLANet_plus
        "use_e2e_wireless_table_rec_model": False,
        # 是否启用无线表格端到端预测模式，不启用时采用表格单元格检测模型预测结果填充至HTML表格，启用时采用端到端表格结构识别模型的单元格预测结果填充至HTML表格。二者在不同场景下效果不同，请根据实际情况选择。
        "use_wireless_table_cells_trans_to_html": False,
        # 是否启用无线表单元格直转HTML模式，不启用时正常通过表格结构识别模块预测得到表格的HTML结构，启用时直接基于无线表单元格检测模型预测的各单元格位置排布与几何关系构建HTML结构。此模式在不同场景下效果不同，请根据实际情况选择。一般来说，建议在无线表单元格检测非常准确的场景下开启此模式，以获得更加准确的端到端表格识别结果。

        "use_table_recognition": True,
        "use_formula_recognition": False,
        "use_region_detection": True,
        "use_seal_recognition": False,
        "use_chart_recognition": False,
    }

    # 合并用户提供的参数
    default_predict_params.update(predict_params)

    # 执行预测
    output = pipeline.predict(**default_predict_params)

    for res in output:
        res.save_to_img(f"./output/{prefix}/")
        res.save_to_html(f"./output/{prefix}/")
        res.save_to_json(f"./output/{prefix}/")
        res.save_to_xlsx(f"./output/{prefix}/")


@timing_decorator(prefix="文档场景信息抽取v3")
def doc_scene_info_extraction(input_path, pipeline_name="PP-ChatOCRv4-doc", prefix=None, **predict_params):
    """
    文档场景信息抽取函数

    Args:
        input_path: 输入文件路径或numpy数组
        pipeline_name: pipeline名称
        prefix: 输出目录前缀
        **predict_params: 预测参数
    """
    # 如果没有提供prefix，则使用pipeline_name作为默认值
    if prefix is None:
        prefix = pipeline_name

    # 获取pipeline实例
    pipeline = get_pipeline_instance(pipeline_name)

    # 默认推理参数
    default_predict_params = {
        "input": input_path,
        "target_size": -1,
        "use_doc_orientation_classify": True,
        "use_common_ocr": True,
        "use_seal_recognition": False,
        "use_table_recognition": True,
        "initial_predictor": False
    }

    # 合并用户提供的参数
    default_predict_params.update(predict_params)

    # 执行预测
    visual_predict_res = pipeline.visual_predict(**default_predict_params)

    visual_info_list = []
    for res in visual_predict_res:
        visual_info_list.append(res["visual_info"])
        layout_parsing_result = res["layout_parsing_result"]
        layout_parsing_result.save_to_img(save_path=f"./output/{prefix}/")
        layout_parsing_result.save_to_json(save_path=f"./output/{prefix}/")
        layout_parsing_result.save_to_html(save_path=f"./output/{prefix}/")
        layout_parsing_result.save_to_xlsx(f"./output/{prefix}/")
        # print(visual_info_list)


@timing_decorator(prefix="公式识别产线")
def formula_recognition(input_path, pipeline_name="formula_recognition", prefix=None, **predict_params):
    """
    公式识别函数

    Args:
        input_path: 输入文件路径或numpy数组
        pipeline_name: pipeline名称
        prefix: 输出目录前缀
        **predict_params: 预测参数
    """
    # 如果没有提供prefix，则使用pipeline_name作为默认值
    if prefix is None:
        prefix = pipeline_name

    # 获取pipeline实例
    pipeline = get_pipeline_instance(pipeline_name)

    # 默认推理参数
    default_predict_params = {
        "input": input_path,
        "use_layout_detection": True,
        "use_doc_orientation_classify": False,
        "use_doc_unwarping": False,
        # "layout_threshold": 0.5,
        "layout_nms": True,
        "layout_unclip_ratio": 1.0,
        "layout_merge_bboxes_mode": "large"
    }

    # 合并用户提供的参数
    default_predict_params.update(predict_params)

    # 执行预测
    output = pipeline.predict(**default_predict_params)

    for res in output:
        res.print()
        res.save_to_img(f"./output/{prefix}/")
        res.save_to_json(f"./output/{prefix}/")


@timing_decorator(prefix="获取{model_name}模型")
def get_model_instance(model_name, **init_params):
    global _all_model_instances

    if model_name not in _all_model_instances:
        # 默认初始化参数
        default_init_params = {
            "device": "cpu",
        }
        # 合并用户提供的参数
        default_init_params.update(init_params)
        from paddlex import create_model
        _all_model_instances[model_name] = create_model(
            model_name=model_name,
            **default_init_params
        )
        logger.info(f"模型创建完成：{model_name}")

    return _all_model_instances[model_name]


@timing_decorator(prefix="{model_name}模型推理")
def model_predict(input_path, model_name="PP-DocLayout-M", prefix=None, **predict_params):
    if prefix is None:
        prefix = model_name

    pipeline = get_model_instance(model_name)

    # 默认推理参数
    default_predict_params = {
        "input": input_path,
        "batch_size": 1,
    }

    # 合并用户提供的参数
    default_predict_params.update(predict_params)

    # 执行预测
    output = pipeline.predict(**default_predict_params)

    for res in output:
        # res.print(json_format=False)
        res.save_to_img(f"./output/{prefix}")
        # res.save_to_json(f"./output/{prefix}")


def mkdirs(target_dir):
    """
       创建目录，如果目录不存在则创建它（包括所有必要的父目录）

       Args:
           target_dir: 目标目录路径，可以是字符串或 pathlib.Path 对象
       """
    from pathlib import Path
    Path(target_dir).mkdir(parents=True, exist_ok=True)
    return target_dir


if __name__ == '__main__':
    """
    测试目的: 使用create_pipeline() 实例化产线对象
    问题：初始化耗时过长

    from paddlex import create_pipeline

    pipeline = create_pipeline(pipeline=[产线名称])
    output = pipeline.predict([输入图片名称])
    for res in output:
        res.print()
        res.save_to_img("./output/")
        res.save_to_json("./output/")

    """
    # clear_directory("../logs/", r=False)
    clear_logs()
    input = TEST_DATA_DIR / "table.jpg"
    input_path = OUTPUT_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章)/page_0.jpg"
    # input_path = OUTPUT_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章)/page_0_image_crop.jpg"
    # input_path = OUTPUT_DATA_DIR  / "25-注会-轻1-财务成本管理[上册](第3章)/page_0_table_crop.jpg"
    cv_img = cv2.imread(TEST_DATA_DIR / "1715339805571/1715339805571_010.jpg")
    # cv2.imwrite("./output/1715339805571_010.jpg", cv_img)
    # cv2.imshow("Detected Regions", cv_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    # ppstructure = PPStructureV3(lang="ch")
    # result = ppstructure.predict(input)
    # for res in result:
    #     res.print(json_format=False)
    #     print(res)

    # 方案2
    # from paddlex import create_model
    # model = create_model(model_name="SLANet")
    # output = model.predict(input=input, batch_size=1)
    # for res in output:
    #     res.print(json_format=False)
    #     # res.save_to_json("./output/res.json")

    """# 方案3 通用版面解析v3"""
    # get_pipeline_instance("PP-StructureV3")
    # general_layout_analysis(input, prefix=f"通用版面解析v3")

    """通用语义分割"""
    # pipeline = create_pipeline(pipeline="semantic_segmentation")
    # output = pipeline.predict(input=input_path, target_size=-1)
    # for res in output:
    #     res.print()
    #     res.save_to_img(save_path="./output/25-注会-轻1-财务成本管理[上册](第3章)/")
    #     res.save_to_json(save_path="./output/25-注会-轻1-财务成本管理[上册](第3章)/")

    # """公式识别产线"""
    # formula_recognition(str(TEST_DATA_DIR / "formula.png"), prefix="公式识别产线")
    #
    # """文档场景信息抽取v3产线"""
    # doc_scene_info_extraction(cv_img, prefix="文档场景信息抽取v3")
    #
    # """通用表格识别"""
    # table_recognition(cv_img, prefix="通用表格识别")

    # vector_info = pipeline.build_vector(visual_info_list, flag_save_bytes_vector=True, retriever_config={
    #     "module_name": "retriever",
    #     "model_name": "embedding-v1",
    #     "base_url": "https://qianfan.baidubce.com/v2",
    #     "api_type": "qianfan",
    #     "api_key": "api_key"  # your api_key
    # })
    # chat_result = pipeline.chat(
    #     key_list=["驾驶室准乘人数"],
    #     visual_info=visual_info_list,
    #     vector_info=vector_info,
    #     chat_bot_config={
    #         "module_name": "chat_bot",
    #         "model_name": "ernie-3.5-8k",
    #         "base_url": "https://qianfan.baidubce.com/v2",
    #         "api_type": "openai",
    #         "api_key": "api_key"  # your api_key
    #     },
    #     retriever_config={
    #         "module_name": "retriever",
    #         "model_name": "embedding-v1",
    #         "base_url": "https://qianfan.baidubce.com/v2",
    #         "api_type": "qianfan",
    #         "api_key": "api_key"  # your api_key
    #     }
    # )
    # print(chat_result)
    """将PDF图片放大2倍进行分析"""

    page_range = (2, 2)
    zoom = 1.5
    input_pdf = TEST_DATA_DIR / "1715339805571.pdf"
    # input_pdf = TEST_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章).pdf"

    filename_without_ext = extract_filename_without_extension(input_pdf)
    filename_without_ext += "_v2"
    # 构建正文pdf
    # 定义检测器函数（带参数）
    def my_roi_detector(page):
        return pdf_get_content_region(page, zoom_factor=zoom, is_show_log=False)


    # crop_pdf_input = OUTPUT_DATA_DIR / "构建正文pdf_mask.pdf"
    # crop_pdf_with_white_mask(input_pdf, crop_pdf_input, my_roi_detector, page_range=page_range)
    #
    # # 获取正文区域
    # # TODO 重构 pdf_get_content_region_v2 与 pdf_get_content_region 移除 zoom_factor参数,在使用时 通过cv2 缩放图片
    # # pages_roi = pdf_get_content_region_v2(crop_pdf_input, page_range=page_range)
    # pages = get_pages_from_range(crop_pdf_input, page_range=page_range)

    pages = get_pages_from_range(input_pdf, page_range=page_range)

    # for page, roi in zip(pages, pages_roi):
    for page in pages:
        # output_dir = mkdirs(TEST_DATA_DIR / "1715339805571_v2")

        cv_img = page_to_img(page, zoom_factor=zoom, color_mode="BGR")
        # show_image_with_roi(cv_img, is_show=True, color_mode="BGR", window_name=f"page_{page.number:02d}")
        # pdf正文区域
        roi = pdf_get_content_region(cv_img, zoom_factor=zoom)
        print(page.number, roi)
        # 只保留ROI区域内容，其余区域填充白色
        cv_img = process_image_roi(cv_img, roi=roi)
        # show_image_with_roi(cv_img, is_show=True, color_mode="BGR", window_name=f"page_{page.number:02d}" )

        # cv2.imwrite(f"{output_dir}/page_{page.number:03d}.jpg", cv_img)
        # @timing_decorator(prefix="单实例模型创建")
        # def layout_detection_test(cv_img):
        #     from paddleocr import LayoutDetection
        #     model = LayoutDetection(model_name="PP-DocLayout_plus-L")
        #     output = model.predict(cv_img, batch_size=1, layout_nms=True)
        #     for res in output:
        #         res.print()
        #         res.save_to_img(save_path=f"./output/PP-DocLayout_plus-L/{page.number}")
        #         res.save_to_json(save_path=f"./output/PP-DocLayout_plus-L/{page.number}")
        # layout_detection_test(cv_img)
        general_layout_analysis(cv_img, prefix=f"{filename_without_ext}/{page.number:02d}")
        # table_recognition(cv_img, prefix=f"通用表格识别/{page.number}")
        # formula_recognition(cv_img, prefix=f"公式识别产线/{page.number}")
        # doc_scene_info_extraction(cv_img, prefix=f"文档场景信息抽取/{page.number}")
        # show_image_with_roi(page, zoom_factor=2.0, is_show=True, color_mode="BGR")

        """独立模型测试"""
        # model_predict(cv_img, model_name="PP-DocLayout_plus-L",prefix=f"PP-DocLayout_plus-L/{page.number:02d}")
        # model_predict(cv_img, model_name="PP-DocLayout-L",prefix=f"PP-DocLayout-L/{page.number:02d}")
        # model_predict(cv_img, model_name="PP-DocBlockLayout")
        # model_predict(cv_img, model_name="PP-DocLayout-M",prefix=f"PP-DocLayout-M/{page.number:02d}")
        # model_predict(cv_img, model_name="PP-DocLayout-S")

    pass
